# -*- coding: utf-8 -*-
"""
Created on Tue Apr 27 10:15:46 2021

@author: Administrator
"""

from __future__ import absolute_import,division,print_function,unicode_literals

import tensorflow as tf
from model import MyModel
import numpy as np
import matplotlib.pyplot as plt
from time import *

mnist = tf.keras.datasets.mnist

# download adn load data
(x_train,y_train),(x_test,y_test) = mnist.load_data()

x_train,x_test = x_train/255.0,x_test/255.0

# imgs = x_test[:3]
# labs = y_test[:3]
# print(labs)
# plot_imgs = np.hstack(imgs)
# plt.imshow(plot_imgs,cmap='gray')
# plt.show()


x_train = x_train[...,tf.newaxis]
x_test = x_test[...,tf.newaxis]

# x_train = x_train[:100]
# y_train = y_train[:100]

#create data generator
train_ds = tf.data.Dataset.from_tensor_slices((x_train,y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test,y_test)).batch(32)

model = MyModel()

#优化器和损失函数
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()
#损失值和准确率
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

#使用tf.GradientTape来训练模型
@tf.function
def train_step(images,labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels,predictions)
    gradients = tape.gradient(loss,model.trainable_variables)
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    
    train_loss(loss)
    train_accuracy(labels,predictions)
    
#测试模型
@tf.function
def test_step(images,labels):
    predictions = model(images)
    t_loss = loss_object(labels,predictions)
    test_loss(t_loss)
    test_accuracy(labels,predictions)
    
EPOCHS = 10
Train_Loss = []
Train_Accuracy = []
Test_Loss = []
Test_Accuracy = []
Epochs = []

for epoch in range(EPOCHS):
    begin_time = time()
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()
    
    for images,labels in train_ds:
        train_step(images,labels)
    for test_images,test_labels in test_ds:
        test_step(test_images,test_labels)
    end_time = time()
    use_time = end_time-begin_time
    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}, use_time {}'
    print(template.format(epoch+1,
                          train_loss.result(),
                          train_accuracy.result()*100,
                          test_loss.result(),
                          test_accuracy.result()*100,
                          use_time))
    # Epochs[epoch] = epoch
    Epochs.append(epoch)
    Train_Loss.append(train_loss.result())
    Train_Accuracy.append(train_accuracy.result())
    Test_Loss.append(test_loss.result())
    Test_Accuracy.append(test_accuracy.result())

# 创建画板
fig = plt.figure(1)

# 创建画纸
ax1 = plt.subplot(2,1,1)
plt.xlabel("EPOCHS")
plt.ylabel("LOSS")
line1 = plt.plot(Epochs,Train_Loss,label="train loss",color='r',marker='o',linestyle='-')
line2 = plt.plot(Epochs,Test_Loss,label="test loss",color='b',marker='s',linestyle='--')
plt.title("Train/Test Loss")
plt.legend()

plt.subplot(2,1,2)
plt.xlabel("EPOCHS")
plt.ylabel("ACCURACY")
line1 = plt.plot(Epochs,Train_Accuracy,label="train accuracy",color='r',marker='o',linestyle='-')
line2 = plt.plot(Epochs,Test_Accuracy,label="test accuracy",color='b',marker='s',linestyle='--')
plt.title("Train/Test Accuracy")
plt.legend()

plt.show()



